package cn.wolfcode.core;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
@Component
@ServerEndpoint("/{token}")
public class WebSocketServer {

    private static final ConcurrentHashMap<String, Session> SESSION_MAP = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(@PathParam("token") String token, Session session) {
        log.info("[WebSocket] 有新的连接：{}", token);
        SESSION_MAP.put(token, session);
    }

    @OnClose
    public void onClose(@PathParam("token") String token) {
        log.info("[WebSocket] 连接关闭：{}", token);
        SESSION_MAP.remove(token);
    }

    @OnError
    public void onError(@PathParam("token") String token, Throwable throwable) {
        SESSION_MAP.remove(token);
        log.error("[WebSocket] 连接出错", throwable);
    }

    public static void sendMessage(String token, String msg) {
        try {
            int count = 0;
            Session session = null;
            do {
                if (count > 4) {
                    log.info("[WebSocket] 重试:{}, 找不到连接：token={}, msg={}", count, token, msg);
                    return;
                }

                session = SESSION_MAP.get(token);
                if (session != null) {
                    log.info("[WebSocket] 获取到连接, 准备发送消息：count={}, token={}", count, token);
                    break;
                }

                TimeUnit.SECONDS.sleep(1);
                count++;
            } while (true);


            session.getBasicRemote().sendText(msg);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
